# 参考 https://github.com/gaoyb7/115drive-webdav/blob/v0.1.8/115/crypto.go
import base64
import os
from typing import Union, List
# noinspection PyPackageRequirements
from Crypto.Cipher import PKCS1_v1_5
# noinspection PyPackageRequirements
from Crypto.PublicKey import RSA


xor_key_seed = bytes([
    0xf0, 0xe5, 0x69, 0xae, 0xbf, 0xdc, 0xbf, 0x8a,
    0x1a, 0x45, 0xe8, 0xbe, 0x7d, 0xa6, 0x73, 0xb8,
    0xde, 0x8f, 0xe7, 0xc4, 0x45, 0xda, 0x86, 0xc4,
    0x9b, 0x64, 0x8b, 0x14, 0x6a, 0xb4, 0xf1, 0xaa,
    0x38, 0x01, 0x35, 0x9e, 0x26, 0x69, 0x2c, 0x86,
    0x00, 0x6b, 0x4f, 0xa5, 0x36, 0x34, 0x62, 0xa6,
    0x2a, 0x96, 0x68, 0x18, 0xf2, 0x4a, 0xfd, 0xbd,
    0x6b, 0x97, 0x8f, 0x4d, 0x8f, 0x89, 0x13, 0xb7,
    0x6c, 0x8e, 0x93, 0xed, 0x0e, 0x0d, 0x48, 0x3e,
    0xd7, 0x2f, 0x88, 0xd8, 0xfe, 0xfe, 0x7e, 0x86,
    0x50, 0x95, 0x4f, 0xd1, 0xeb, 0x83, 0x26, 0x34,
    0xdb, 0x66, 0x7b, 0x9c, 0x7e, 0x9d, 0x7a, 0x81,
    0x32, 0xea, 0xb6, 0x33, 0xde, 0x3a, 0xa9, 0x59,
    0x34, 0x66, 0x3b, 0xaa, 0xba, 0x81, 0x60, 0x48,
    0xb9, 0xd5, 0x81, 0x9c, 0xf8, 0x6c, 0x84, 0x77,
    0xff, 0x54, 0x78, 0x26, 0x5f, 0xbe, 0xe8, 0x1e,
    0x36, 0x9f, 0x34, 0x80, 0x5c, 0x45, 0x2c, 0x9b,
    0x76, 0xd5, 0x1b, 0x8f, 0xcc, 0xc3, 0xb8, 0xf5
])

xor_client_key = bytes([
    0x78, 0x06, 0xad, 0x4c, 0x33, 0x86, 0x5d, 0x18,
    0x4c, 0x01, 0x3f, 0x46
])

rsa_public_key = bytes("-----BEGIN PUBLIC KEY-----\n" +
                       "MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQCGhpgMD1okxLnUMCDNLCJwP/P0\n" +
                       "UHVlKQWLHPiPCbhgITZHcZim4mgxSWWb0SLDNZL9ta1HlErR6k02xrFyqtYzjDu2\n" +
                       "rGInUC0BCZOsln0a7wDwyOA43i5NO8LsNory6fEKbx7aT3Ji8TZCDAfDMbhxvxOf\n" +
                       "dPMBDjxP5X3zr7cWgwIDAQAB\n" +
                       "-----END PUBLIC KEY-----", 'utf-8')

rsa_server_key = RSA.importKey(rsa_public_key)


def generate_key() -> bytes:
    return os.urandom(16)


def compare(a: List[int], b: str) -> bool:
    seq = b.replace('[', '').replace(']', '').split(' ')
    return a == [int(x) for x in seq]


def pickcode2bytes(pc: str) -> bytes:
    s = f'{{"pickcode":"{pc}"}}'
    return bytes(s, 'utf-8')


def print_bytes(bts: Union[bytes, bytearray]):
    print([int(x) for x in bts])


def xor_transform(data: bytearray, key: bytes):
    data_size, key_size = len(data), len(key)
    mod = data_size % 4
    if mod > 0:
        for i in range(mod):
            data[i] ^= key[i % key_size]
    for i in range(mod, data_size):
        data[i] ^= key[(i - mod) % key_size]
    return data


def xor_derive_key(seed: bytes, size: int):
    key = bytearray(size)
    for i in range(size):
        key[i] = (seed[i] + xor_key_seed[size * i]) & 0xff
        key[i] ^= xor_key_seed[size * (size - i - 1)]
    return key


def rsa_decrypt(input_data: bytes):
    output = bytearray()
    cipher_size, block_size = len(input_data), rsa_server_key.size_in_bytes()
    for offset in range(0, cipher_size, block_size):
        slice_size = block_size
        if offset + slice_size > cipher_size:
            slice_size = cipher_size - offset

        n = int.from_bytes(input_data[offset: offset + slice_size], byteorder='big')
        m = pow(n, rsa_server_key.e, rsa_server_key.n)
        b = m.to_bytes((m.bit_length() + 7) // 8, byteorder='big')
        index = b.find(b'\x00', 1)
        if index < 0:
            return None
        output += b[index + 1:]

    return output


def encode(input_data: bytes, key: bytes):
    buffer = bytearray(16 + len(input_data))
    buffer[:16] = key
    buffer[16:] = xor_transform(bytearray(input_data), xor_derive_key(key, 4))
    buffer[16:] = reversed(buffer[16:])
    buffer[16:] = xor_transform(buffer[16:], xor_client_key)
    cipher = PKCS1_v1_5.new(rsa_server_key)
    encrypted_data = cipher.encrypt(bytes(buffer))
    return base64.b64encode(encrypted_data).decode()


def decode(input_data: str, key: bytes):
    encrypted_data = base64.b64decode(input_data.encode())
    decrypted_data = rsa_decrypt(encrypted_data)
    output = bytearray(len(decrypted_data) - 16)
    output[:] = decrypted_data[16:]
    output = xor_transform(output, xor_derive_key(decrypted_data[:16], 12))
    output[:] = reversed(output)
    output = xor_transform(output, xor_derive_key(key, 4))
    return bytes(output)


if __name__ == '__main__':
    pass
